import cv2
import numpy as np
import setting


def find_pic_old(large_image: str, small_image: str, debug: bool = False):
    method = cv2.TM_SQDIFF_NORMED
    # method = cv2.TM_CCOEFF_NORMED

    # Read the images from the file
    large_image = cv2.imread(r'C:\Users\Hanse\Documents\leidian\Pictures\Screenshots\capture.png')
    small_image = cv2.imread(r'assets\return_sea.png')
    # large_image = cv2.imread('assets/Screenshot_20210426-181401.png')

    result = cv2.matchTemplate(small_image, large_image, method)

    # We want the minimum squared difference
    mn, _, mnLoc, _ = cv2.minMaxLoc(result)

    # Draw the rectangle:
    # Extract the coordinates of our best match
    MPx, MPy = mnLoc

    # Step 2: Get the size of the template. This is the same size as the match.
    trows, tcols = small_image.shape[:2]

    # Step 3: Draw the rectangle on large_image
    cv2.rectangle(large_image, (MPx, MPy), (MPx + tcols, MPy + trows), (0, 0, 255), 2)

    # Display the original image with the rectangle around the match.
    if debug:
        # 打印坐标
        print(MPx, MPy, MPx + tcols, MPy + trows, mn)
        cv2.imshow('output', large_image)
        cv2.waitKey(0)

    return MPx, MPy, MPx + tcols, MPy + trows, int(MPx + tcols / 2), int(MPy + trows / 2)


def find_pic_ex(large_image: str, small_image: str, threshold: float = setting.img_threshold, debug: bool = False):
    """
    是否找到图片，含坐标

    :param large_image:
    :param small_image:
    :param threshold:
    :param debug:
    :return:
    """
    img_rgb = cv2.imread(large_image)
    template = cv2.imread(small_image)
    h, w = template.shape[:-1]

    res = cv2.matchTemplate(img_rgb, template, cv2.TM_CCOEFF_NORMED)
    # threshold = 0.8
    loc = np.where(res >= threshold)
    for pt in zip(*loc[::-1]):  # Switch collumns and rows
        if not debug:
            return pt[0], pt[1], pt[0] + w, pt[1] + h, int(pt[0] + w / 2), int(pt[1] + h / 2)
        print(pt[0], pt[1], pt[0] + w, pt[1] + h, int(pt[0] + w / 2), int(pt[1] + h / 2))
        cv2.rectangle(img_rgb, pt, (pt[0] + w, pt[1] + h), (0, 0, 255), 2)

    if debug:
        cv2.imwrite('result.png', img_rgb)
    return -1, -1, -1, -1, -1, -1


def find_pic(large_image: str, small_image: str, threshold: float = setting.img_threshold, debug: bool = False):
    """
    是否找到图片，不含坐标

    :param large_image:
    :param small_image:
    :param threshold:
    :param debug:
    :return:
    """
    x1, y1, x2, y2, x, y = find_pic_ex(large_image, small_image, threshold, debug)
    if x1 == y1 == x2 == y2 == x == y == -1:
        return False
    else:
        return True


if __name__ == '__main__':
    large_image = r'C:\Users\Hanse\Documents\leidian\Pictures\Screenshots\capture.png'
    # small_image = 'assets/overview.png'
    small_image = 'assets/return_sea.png'
    # find_pic_old(large_image, 'assets/overview.png', True)

    print(find_pic_ex(large_image, small_image, 0.8, True))
    # print(find_pic_old(large_image, small_image, True))
